import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from .resnet import resnet34


class CNN_encoder(nn.Module):
    def __init__(self):
        super(CNN_encoder, self).__init__()
        self.backbone = resnet34()  # resnet 将图片从 h x w 变为 h/32 x w/16的特征图
        self.context = nn.LSTM(512, 512, 2, batch_first=True)

    def forward(self, x):
        encoder_features = self.backbone(x)
        encoder_features = encoder_features.squeeze(2).permute(0, 2, 1)  # b,c,1,l -> b,c,l -> b,l,c
        encoder_features, _ = self.context(encoder_features)  # b,l,c
        return encoder_features


class CRNN(nn.Module):
    def __init__(self, num_class):
        super(CRNN, self).__init__()
        self.encoder = CNN_encoder()
        self.dropout = nn.Dropout(0.3)  # 因为手写数据集比较小，因此需要考虑避免模型过拟合的问题
        self.cls_op = nn.Sequential(nn.Linear(512, 1024),
                                    nn.ReLU(True),
                                    nn.Linear(1024, num_class))

    def forward(self, x):
        features = self.encoder(x)
        features = self.dropout(features)
        logits = self.cls_op(features)
        # logits = F.log_softmax(logits, dim=-1)
        logits = F.log_softmax(logits, dim=2)
        return logits


class Residual_block(nn.Module):
    def __init__(self, c_in, c_out, stride):
        super(Residual_block, self).__init__()
        self.downsample = None
        flag = False
        if isinstance(stride, tuple):
            if stride[0] > 1:
                self.downsample = nn.Sequential(nn.Conv2d(c_in, c_out, 3, stride, 1),
                                                nn.BatchNorm2d(c_out, momentum=0.01))
                flag = True
        else:
            if stride > 1:
                self.downsample = nn.Sequential(nn.Conv2d(c_in, c_out, 3, stride, 1),
                                                nn.BatchNorm2d(c_out, momentum=0.01))
                flag = True
        if flag:
            self.conv1 = nn.Sequential(nn.Conv2d(c_in, c_out, 3, stride, 1),
                                       nn.BatchNorm2d(c_out, momentum=0.01))
        else:
            self.conv1 = nn.Sequential(nn.Conv2d(c_in, c_out, 1, stride, 0),
                                       nn.BatchNorm2d(c_out, momentum=0.01))
        self.conv2 = nn.Sequential(nn.Conv2d(c_out, c_out, 3, 1, 1),
                                   nn.BatchNorm2d(c_out, momentum=0.01))
        self.relu = nn.ReLU()

    def forward(self, x):
        residual = x
        conv1 = self.conv1(x)
        conv2 = self.conv2(conv1)
        if self.downsample is not None:
            residual = self.downsample(residual)
        return self.relu(residual + conv2)


class CNN_ResNet18_for_cnc(nn.Module):
    def __init__(self, c_in):
        super(CNN_ResNet18_for_cnc, self).__init__()
        self.block0 = nn.Sequential(nn.Conv2d(c_in, 64, 7, 1, 1), nn.BatchNorm2d(64))  # 48*480 -> 44*476
        self.block1 = self._make_layer(64, 128, (2, 2), 1)  # -> 22*238
        self.block2 = self._make_layer(128, 256, (2, 2), 1)  # -> 11*119
        self.block3 = self._make_layer(256, 512, (2, 2), 1)  # -> 6*60
        self.block4 = self._make_layer(512, 512, (2, 1), 1)  # -> 3*60
        self.block5 = self._make_layer(512, 512, (3, 1), 0)  # -> 1*60

    def _make_layer(self, c_in, c_out, stride, repeat=2):
        layers = []
        layers.append(Residual_block(c_in, c_out, stride))
        for i in range(repeat):
            layers.append(Residual_block(c_out, c_out, 1))
        return nn.Sequential(*layers)

    def forward(self, x):
        block0 = self.block0(x)
        block1 = self.block1(block0)
        block2 = self.block2(block1)
        block3 = self.block3(block2)
        block4 = self.block4(block3)
        block5 = self.block5(block4)
        block5 = block5.squeeze(2).permute(0, 2, 1).contiguous()
        return block5


class CRNN_OCR_for_cnc(nn.Module):
    def __init__(self, in_channel, num_classes):
        super(CRNN_OCR_for_cnc, self).__init__()
        self.CNN_encoder = CNN_ResNet18_for_cnc(c_in=in_channel)
        self.cls_op = nn.Linear(512, num_classes)

    def forward(self, inputs):
        features = self.CNN_encoder(inputs)
        logits = self.cls_op(features)
        logits = F.log_softmax(logits, dim=2)
        return logits  # bs x 60 x nclass
